import warnings

import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule, xavier_init  # todo:mmcv update
from mmcv.cnn import ConvModule
from mmcv.runner import BaseModule, auto_fp16

from ..builder import NECKS
import pdb


@NECKS.register_module()
class FPN_NIL_slim_deconv(BaseModule):
    r"""Feature Pyramid Network. (with num ins laterals)

    This is an implementation of paper `Feature Pyramid Networks for Object
    Detection <https://arxiv.org/abs/1612.03144>`_.

    wxz: replace interpolate function by deconv layer

    Args:
        in_channels (List[int]): Number of input channels per scale.
        out_channels (int): Number of output channels (used at each scale)
        num_outs (int): Number of output scales.
        start_level (int): Index of the start input backbone level used to
            build the feature pyramid. Default: 0.
        end_level (int): Index of the end input backbone level (exclusive) to
            build the feature pyramid. Default: -1, which means the last level.
        add_extra_convs (bool | str): If bool, it decides whether to add conv
            layers on top of the original feature maps. Default to False.
            If True, its actual mode is specified by `extra_convs_on_inputs`.
            If str, it specifies the source feature map of the extra convs.
            Only the following options are allowed

            - 'on_input': Last feat map of neck inputs (i.e. backbone feature).
            - 'on_lateral':  Last feature map after lateral convs.
            - 'on_output': The last output feature map after fpn convs.
        extra_convs_on_inputs (bool, deprecated): Whether to apply extra convs
            on the original feature from the backbone. If True,
            it is equivalent to `add_extra_convs='on_input'`. If False, it is
            equivalent to set `add_extra_convs='on_output'`. Default to True.
        relu_before_extra_convs (bool): Whether to apply relu before the extra
            conv. Default: False.
        no_norm_on_lateral (bool): Whether to apply norm on lateral.
            Default: False.
        conv_cfg (dict): Config dict for convolution layer. Default: None.
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (str): Config dict for activation layer in ConvModule.
            Default: None.
        upsample_cfg (dict): Config dict for interpolate layer.
            Default: `dict(mode='nearest')`
        init_cfg (dict or list[dict], optional): Initialization config dict.
        method (str): method to fuse features in FPN.

            origin: standard FPN process.
                    featmap[3]→→→ ●
                                  ↓
                    featmap[2]→→→ + →→→ out[2]
                                  ↓
                    featmap[1]→→→ + →→→ out[1]
                                  ↓
                    featmap[0]→→→ + →→→ out[0]

            method1:
                    featmap[3]→→→ ● →→→ out[2]
                                  ↓
                    featmap[2]→→→ + →→→ out[1]
                                  ↓
                    featmap[1]→→→ + →→→ out[0]
                                  ↑
                    featmap[0]→→→ ●

            method2:
                    featmap[3]→→→ ●
                                  ↓
                    featmap[2]→→→ + →→→→→→→→→→→→→→→→→→→→→ out[2]
                                  ↓
                    featmap[1]→→→ + →→→→→→→→→→→→→→→→→→→→→ out[1]
                                  ↓
                    featmap[0]→→→ ● →→→(max_pooling 2)→→→ out[0]

    Example:
        >>> import torch
        >>> in_channels = [2, 3, 5, 7]
        >>> scales = [340, 170, 84, 43]
        >>> inputs = [torch.rand(1, c, s, s)
        ...           for c, s in zip(in_channels, scales)]
        >>> self = FPN(in_channels, 11, len(in_channels)).eval()
        >>> outputs = self.forward(inputs)
        >>> for i in range(len(outputs)):
        ...     print(f'outputs[{i}].shape = {outputs[i].shape}')
        outputs[0].shape = torch.Size([1, 11, 340, 340])
        outputs[1].shape = torch.Size([1, 11, 170, 170])
        outputs[2].shape = torch.Size([1, 11, 84, 84])
        outputs[3].shape = torch.Size([1, 11, 43, 43])
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        num_outs,
        start_level=0,
        end_level=-1,
        add_extra_convs=False,
        extra_convs_on_inputs=True,
        relu_before_extra_convs=False,
        no_norm_on_lateral=False,
        conv_cfg=None,
        norm_cfg=None,
        act_cfg=None,
        upsample_cfg=dict(mode="nearest"),
        init_cfg=dict(type="Xavier", layer="Conv2d", distribution="uniform"),
        deconv_with_relu=True,
        deconv_with_1x1conv=False,
        deconv_param=[4, 2, 1],
        method="origin",  # "origin","method1","method2"
    ):
        super(FPN_NIL_slim_deconv, self).__init__(init_cfg)
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.relu_before_extra_convs = relu_before_extra_convs
        self.no_norm_on_lateral = no_norm_on_lateral
        self.fp16_enabled = False
        self.upsample_cfg = upsample_cfg.copy()
        self.method = method

        if end_level == -1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level < inputs, no extra level is allowed
            self.backbone_end_level = end_level
            assert end_level <= len(in_channels)
            assert num_outs == end_level - start_level
        self.start_level = start_level
        self.end_level = end_level
        self.add_extra_convs = add_extra_convs
        assert isinstance(add_extra_convs, (str, bool))
        if isinstance(add_extra_convs, str):
            # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output'
            assert add_extra_convs in ("on_input", "on_lateral", "on_output")
        elif add_extra_convs:  # True
            if extra_convs_on_inputs:
                # TODO: deprecate `extra_convs_on_inputs`
                warnings.simplefilter("once")
                warnings.warn(
                    '"extra_convs_on_inputs" will be deprecated in v2.9.0,'
                    'Please use "add_extra_convs"',
                    DeprecationWarning,
                )
                self.add_extra_convs = "on_input"
            else:
                self.add_extra_convs = "on_output"

        self.lateral_convs = nn.ModuleList()
        #  self.fpn_convs = nn.ModuleList()
        self.de_convs = nn.ModuleList()
        # num ins lateral layers, but num outs fpn layers
        for i in range(self.start_level, self.num_ins):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,
                act_cfg=act_cfg,
                inplace=False,
            )
            self.lateral_convs.append(l_conv)
            #  de_conv = nn.ConvTranspose2d(in_channels[i],out_channels,3,1,1,0)

            # deconv  can use 1x1 to reduce computation
            deconv_feat_channels = 16 if deconv_with_1x1conv else out_channels
            if deconv_with_relu:
                de_conv = [
                    nn.ConvTranspose2d(
                        out_channels,
                        deconv_feat_channels,
                        deconv_param[0],
                        deconv_param[1],
                        deconv_param[2],
                    ),
                    nn.BatchNorm2d(deconv_feat_channels),
                    nn.ReLU(),
                ]
            else:
                de_conv = [
                    nn.ConvTranspose2d(
                        out_channels,
                        deconv_feat_channels,
                        deconv_param[0],
                        deconv_param[1],
                        deconv_param[2],
                    ),
                    nn.BatchNorm2d(deconv_feat_channels),
                ]
            if deconv_with_1x1conv:
                de_1x1_conv = [
                    nn.Conv2d(deconv_feat_channels, out_channels, 1, 1),
                    nn.BatchNorm2d(out_channels),
                    nn.ReLU(),
                ]
            else:
                de_1x1_conv = []
            de_conv.extend(de_1x1_conv)
            de_conv = nn.Sequential(*de_conv)

            self.de_convs.append(de_conv)

    # mmcv update
    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        """Initialize the weights of FPN module."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution="uniform")

    @auto_fp16()
    def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]
        # import pdb; pdb.set_trace()

        # build top-down path
        used_backbone_levels = len(laterals)
        out = [0 for i in range(self.start_level, self.end_level)]

        if self.method == "origin":
            for i in range(used_backbone_levels - 1, 0, -1):
                laterals[i - 1] += self.de_convs[i](laterals[i])
                out[i - 1] = laterals[i - 1]
        elif self.method == "method1":
            out[-1] = laterals[-1]
            for i in range(used_backbone_levels - 1, 1, -1):
                laterals[i - 1] += self.de_convs[i](laterals[i])
                out[i - 2] = laterals[i - 1]
            out[0] += nn.MaxPool2d(3, stride=2, padding=1)(laterals[0])
        elif self.method == "method2":
            for i in range(used_backbone_levels - 1, 0, -1):
                laterals[i - 1] += self.de_convs[i](laterals[i])
                out[i - 1] = laterals[i - 1]
            out[0] = nn.MaxPool2d(3, stride=2, padding=1)(out[0])
        else:
            raise ValueError(f"{self.method} is not supported.")

        return tuple(out)
